from datetime import datetime
import pandas as pd
import backtrader as bt
import akshare as ak
import os

# 导入自己编写的策略
from strategy.strategy_turtle import CTurtleStrategy
from strategy.strategy_turtle import CTurtleSizer

from strategy.strategy_dual_ma import CDualMAStrategy

from strategy.strategy_simple_ma import CSimpleMAStrategy

from strategy.strategy_bollinger_bands import CBollingerBandsStrategy

from strategy.strategy_demo import CDemoStrategy

class CQuantMainFrame():
    code = "sz002415"                        # 股票代码
    init_cash = 100000.0                      # 初始资金
    commission = 0.02                        # 交易手续费
    start_date = datetime(2020, 1, 1)
    end_date = datetime(2024, 7, 31)

    def __init__(self, **kwargs):
        # 设置数据存储目录
        self.data_dir = os.getcwd() + "/../../data/"
        if not os.path.isdir(self.data_dir):
            os.mkdir(self.data_dir)

        print(f"data directory: {self.data_dir}")

        # 获取回测数据
        datasource = kwargs.get("datasource", "net")
        if datasource == "net":
            self.__load_netdata()
        else:
            self.__load_csvdata()

        self.bt_data = bt.feeds.PandasData(dataname=self.stock_data_df,
                                           fromdate=self.start_date,
                                           todate=self.end_date)
    def __load_csvdata(self):
        self.stock_data_df = pd.read_csv(f"{self.data_dir}/{self.code}.csv", parse_dates=True, index_col=0)
        pass
    
    def __load_netdata(self):
        stock_df = ak.stock_zh_a_daily(symbol=self.code, 
                                       start_date=self.start_date.strftime('%Y%m%d'), 
                                       end_date=self.end_date.strftime('%Y%m%d'))
                                    #adjust='hfq')

        stock_df.to_csv(f"{self.data_dir}/{self.code}.csv", encoding="utf-8-sig", index=False)
        stock_df.set_index(stock_df['date'],inplace=True)
        stock_df.drop('date',axis=1,inplace=True)
        stock_df.index = pd.to_datetime(stock_df.index)
        self.stock_data_df = stock_df
        stock_df.to_csv()
    
    def run_turtle_strategy(self):

        # 初始化回测系统
        cerebro = bt.Cerebro()
        # 将数据传入回测系统
        cerebro.adddata(self.bt_data)

        # 将交易策略加入回测系统
        cerebro.addstrategy(CTurtleStrategy)

        # 设置初始资本和交易手续费
        cerebro.broker.setcash(self.init_cash)
        cerebro.broker.setcommission(commission=self.commission)
        
        # 将仓位控制加入回测系统
        cerebro.addsizer(CTurtleSizer)

        # 将分析器加入回测系统
        cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name="SharpeRatio")
        cerebro.addanalyzer(bt.analyzers.DrawDown, _name="DrawDown")
        cerebro.addanalyzer(bt.analyzers.TradeAnalyzer, _name="TradeAnalyzer")

        # 运行回测系统
        results = cerebro.run(maxcpus=4)
        strat = results[0]
        print(f'最终资金: {round(cerebro.broker.getvalue(),2)}') # 获取回测结束后的总资金
        print('夏普比率:', strat.analyzers.SharpeRatio.get_analysis())
        print('回撤指标:', strat.analyzers.DrawDown.get_analysis())
        print('交易信息:', strat.analyzers.TradeAnalyzer.get_analysis())

        for k, v in strat.analyzers.TradeAnalyzer.get_analysis().items():
                print(k, '->', v)
        
        cerebro.plot(style='candle', volume=True, voloverlay=True)

    def run_dual_ma_strategy(self):
        cerebro = bt.Cerebro()
        cerebro.addstrategy(CDualMAStrategy)

        cerebro.broker.setcash(self.init_cash)
        cerebro.broker.setcommission(commission=self.commission)

        # 购买固定的股票数
        cerebro.addsizer(bt.sizers.FixedSize, stake=100)    
        cerebro.adddata(self.bt_data)
        cerebro.run(maxcpus=4)

        # 当然也可以选自己喜欢的主题色
        mycolors = ['#729ece', '#ff9e4a', '#67bf5c', 
                '#ed665d', '#ad8bc9', '#a8786e', 
                '#ed97ca', '#a2a2a2', '#cdcc5d', '#6dccda']
        cerebro.plot(style='candel',    # 设置主图行情数据的样式为蜡烛图
                     lcolors=mycolors,  # 重新设置主题颜色
                     plotdist=0.1,      # 设置图形之间的间距
                     numfigs=1, # 是否将图形拆分成多幅图展示，如果时间区间比较长，建议分多幅展示
                     barup = '#ff9896', bardown='#98df8a', # 设置蜡烛图上涨和下跌的颜色
                     volup='#ff9896', voldown='#98df8a', # 设置成交量在行情上涨和下跌情况下的颜色
                     )

    def run_simple_ma_strategy(self):
        cerebro = bt.Cerebro()
        cerebro.addstrategy(CSimpleMAStrategy)

        cerebro.broker.setcash(self.init_cash)
        cerebro.broker.setcommission(commission=self.commission)

        # 购买固定的股票数
        cerebro.addsizer(bt.sizers.FixedSize, stake=100)    
        cerebro.adddata(self.bt_data, name=self.code)
        cerebro.run(maxcpus=4)

        cerebro.plot(style='bar',
                     barup = "red",
                     bardown = "green",
                     volume = True,
                     voloverlay = False
                     )
        
    def run_bollinger_bands_strategy(self):
        
        # 创建Cerebro引擎
        cerebro = bt.Cerebro()
        
        # 设置初始资金
        cerebro.broker.setcash(100000.0)
        
       
        # 将数据添加到Cerebro引擎中
        cerebro.adddata(self.bt_data, name=self.code)
        
        # 添加MACD策略
        cerebro.addstrategy(CBollingerBandsStrategy)
        
        # 设置佣金为0.1%
        cerebro.broker.setcommission(commission=0.001)
        
        # 添加分析指标
        # cerebro.addanalyzer(bt.analyzers.Returns, _name='returns')
        # cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name='sharpe')
        # cerebro.addanalyzer(bt.analyzers.DrawDown, _name='drawdown')
        
        # 运行回测
        print('Starting Portfolio Value: %.2f' % cerebro.broker.getvalue())
        results = cerebro.run()
        print('Final Portfolio Value: %.2f' % cerebro.broker.getvalue())
        
        # 获取回测结果
        # strat = results[0]
        # returns = strat.analyzers.returns.get_analysis()
        # sharpe = strat.analyzers.sharpe.get_analysis()
        # drawdown = strat.analyzers.drawdown.get_analysis()
        
        # 打印回测指标
        # print('Annualized Return: %.2f%%' % (returns['rnorm100']))
        # print('Sharpe Ratio: %.2f' % (sharpe['sharperatio']))
        # print('Max Drawdown: %.2f%%' % (drawdown['max']['drawdown']))
        # print('Max Drawdown Period: %s' % (drawdown['max']['len']))
        
        
        # 绘制回测结果
        cerebro.plot(style='candel')


    def run_demo_strategy(self):
        # 创建了一个机器人大脑（Cerebro），同时隐含创建了一个borker（券商)
        cerebro = bt.Cerebro()

        # 设置初始资金（向券商里面存入资金：入金）
        cerebro.broker.setcash(100000.0)

        cerebro.broker.setcommission(commission=0.000485)

        # 将数据加入回测系统（可以同时加入多条数据）
        cerebro.adddata(self.bt_data)
        #cerebro.adddata(self.bt_data)

        ## 添加策略，并给策略传递参数
        #cerebro.addstrategy(CDemoStrategy, myparam=20, exitbars=4)

        # 增加多参数的策略
        strats = cerebro.optstrategy(CDemoStrategy, maperiod=range(10, 31))


        # A股一手100股
        cerebro.addsizer(bt.sizers.FixedSize, stake=100)

        # 显示了机器人在券商那里存有多少钱
        print('Starting Portfolio Value: %.2f' % cerebro.broker.getvalue())
        # 让机器人大脑开始运行。
        cerebro.run()
        print('Final Portfolio Value: %.2f' % cerebro.broker.getvalue())

if __name__ == "__main__":
    main_frame = CQuantMainFrame(datasource="csv")
    #main_frame.run_turtle_strategy()
    #main_frame.run_dual_ma_strategy()
    #main_frame.run_simple_ma_strategy()
    #main_frame.run_bollinger_bands_strategy()
    main_frame.run_demo_strategy()